fix(chunk-grids): make array creation O(1) in per-dimension chunk count - #4218
fix(chunk-grids): make array creation O(1) in per-dimension chunk count#4218d-v-b wants to merge 16 commits into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #4218 +/- ##
==========================================
+ Coverage 94.21% 94.29% +0.07%
==========================================
Files 92 92
Lines 12863 12898 +35
==========================================
+ Hits 12119 12162 +43
+ Misses 744 736 -8
🚀 New features to boost your workflow:
|
Merging this PR will improve performance by 11.82%
Performance Changes
Tip Curious why this is faster? Comment Comparing Footnotes
|
…k_grid_metadata The defensive TypeError branch is unreachable through the public API, so exercise it directly with a stub dimension object. This was the only genuinely uncovered patch line in zarr-developers#4218; the other lines codecov flagged came from a partial coverage upload. Assisted-by: ClaudeCode:claude-fable-5
Chunk normalization now returns a ChunkGrid whose uniform dimensions are stored as FixedDimension (size + extent) instead of one array entry per chunk, so create_array(shape=(2**62,), chunks=(1,)) succeeds instantly instead of raising "array is too big", and chunks=(1, 1) on a (2**31, 2**31) array no longer allocates ~17 GB per dimension. Explicit per-chunk lists collapse to FixedDimension when they describe a regular grid; genuinely irregular lists become VaryingDimension. Both variants bind chunk sizes to their extent, so the two forms carry the same invariants. The intermediate ChunksTuple type and as_regular_shape helper are removed; create_chunk_grid_metadata consumes the grid directly and serializes uniform dimensions of mixed rectilinear grids as the spec's bare-int step-size shorthand. Creation-time counterpart of the zarr-developersgh-4174 indexing fix. Assisted-by: ClaudeCode:claude-fable-5
Since normalize_chunks_nd returns ChunkGrid dimensions instead of int64 arrays, is_regular_1d only ever receives plain Python sequences; the vectorized numpy path was unreachable in production code. Remove it and narrow the signatures of is_regular_1d / is_regular_nd to Sequence[int]. Assisted-by: ClaudeCode:claude-fable-5
…k_grid_metadata The defensive TypeError branch is unreachable through the public API, so exercise it directly with a stub dimension object. This was the only genuinely uncovered patch line in zarr-developers#4218; the other lines codecov flagged came from a partial coverage upload. Assisted-by: ClaudeCode:claude-fable-5
…ersion The stateful hypothesis tests convert generated chunk grid metadata back into a create_array chunks= argument. Bare-int dimensions — the spec's step-size shorthand, now produced when a uniform dimension of a mixed grid collapses — were wrapped as single-element lists, turning "repeat to cover the axis" into "exactly one chunk" and failing the sum-to-span check (e.g. chunks=[1] for span 3). Extract the conversion into chunks_param_from_rectilinear, pass bare ints through unchanged, and widen ChunksLike to admit mixed int | sequence per-dimension specs, which the normalizer already accepted. Assisted-by: ClaudeCode:claude-fable-5
a6d27c7 to
72af8cb
Compare
|
🤖 AI text below 🤖 Rebased onto current RebaseThe conflict was with #4257 ("accept numpy integers as chunk sizes"), which landed after this branch was cut and rewrote the same region of
New commit: explicit per-chunk lists stay rectilinear (fixes #4272)The explicit-list branch of Relation to #4290#4290 fixes the same issue on current |
Drop the regular-grid collapse for explicit list input in normalize_chunks_1d: a per-chunk size list now always produces VaryingDimension, even when the sizes are uniform or uniform plus a short tail. Scalar specs (ints, numpy integers, and the -1 sentinel) still produce FixedDimension, which is the path that makes array creation O(1) in per-dimension chunk count; explicit lists are already O(n) in the input, so nothing is lost. The grid kind now follows the input syntax, matching 3.2.x behavior: previously an explicitly rectilinear spec whose edges happened to look regular was silently stored as RegularChunkGridMetadata, which changes resize semantics — a regular grid grows by extending the uniform pattern while a rectilinear grid appends an edge chunk, breaking append-oriented layouts like (168,) * 13 + (24,). This is resolution option 1 from zarr-developersgh-4272. Uniform dimensions of mixed scalar/list specs still serialize as the spec's bare-int step-size shorthand; explicit lists serialize as edge lists. The touched-chunk-keys assertions in the resize regression test follow the approach from zarr-developersgh-4290, which fixes the same issue on main via a requested_rectilinear flag. Fixes zarr-developers#4272 Co-authored-by: Shurong Cao <CAOShurong@users.noreply.github.com> Assisted-by: ClaudeCode:claude-fable-5
72af8cb to
535fa0a
Compare
…me grid The stateful hypothesis tests in CI failed in from_array on a rectilinear source whose edges happen to be uniform (e.g. chunk_shapes=((3,),)): _parse_keep_array_attr and _info guarded data.chunks / data.shards with the runtime ChunkGrid.is_regular, but the runtime grid collapses uniform rectilinear dimensions to FixedDimension as an optimization, while metadata.chunks raises for any array whose stored grid is rectilinear. Before explicit per-chunk lists stopped collapsing at creation, the two notions of regularity always agreed for created arrays, so the mismatch was unreachable. Add _stored_chunk_grid_is_regular, which dispatches on the metadata kind, and use it at the three sites that read .chunks/.shards. A uniform-edged rectilinear source now round-trips through from_array with its grid kind intact. Assisted-by: ClaudeCode:claude-fable-5
|
@maxrjones could you review this when you get a chance? It fixes two issues and so I consider it high priority for our next release |
maxrjones
left a comment
There was a problem hiding this comment.
So much better ! Thanks @d-v-b! Also sorry for my lack of communication about a review timeline.
On your unpleasant surprises question, below are three tests that I think need to pass before the PR is merged. Everything else I came up with are nits that can go in follow-ups.
Test 1 — the default from_array must work on a rectilinear source (currently: NotImplementedError on .chunks from the data-copy path; the PR's own test hides this behind write_data=False):
def test_from_array_rectilinear_default_copies_data() -> None:
src = zarr.create_array(MemoryStore(), shape=(6,), chunks=[[2, 2, 2]], dtype="uint8")
src[:] = 7
dst = zarr.from_array(MemoryStore(), data=src, name="0") # write_data=True default
assert isinstance(dst.metadata, ArrayV3Metadata)
assert dst.metadata.chunk_grid == src.metadata.chunk_grid
assert (dst[:] == src[:]).all()Test 2 — chunks="keep" must preserve the bare-int shorthand (currently fails with assert ((3, 3, 3), (10, 20)) == (3, (10, 20)) — the shorthand is expanded per-chunk, which is the O(nchunks) path and changes the copy's resize semantics):
def test_from_array_keep_preserves_bare_int_shorthand() -> None:
src = zarr.create_array(MemoryStore(), shape=(9, 30), chunks=(3, [10, 20]), dtype="uint8")
dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=False)
grid = dst.metadata.chunk_grid
assert isinstance(grid, RectilinearChunkGridMetadata)
assert grid.chunk_shapes == (3, (10, 20))Deliberately small: a shape=(2**62, 30) variant would pin the O(1) property too, but before the fix it hangs rather than fails, which makes a bad red test. After fixing, consider using shape=(2**62, 30)
Test 3 — rectilinear-sharded arrays: .info works and sharding survives from_array (currently fails at src.info with NotImplementedError on .shards; the tail assertions then catch the silent sharding drop):
def test_from_array_keep_preserves_rectilinear_sharding() -> None:
src = zarr.create_array(
MemoryStore(), shape=(100,), chunks=(10,), shards=[[50, 50]], dtype="int32"
)
assert src.info is not None # _info must not raise on the shards accessor
dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=False)
assert any(isinstance(c, ShardingCodec) for c in dst.metadata.codecs)
assert dst.metadata.chunk_grid == src.metadata.chunk_gridAnother option is to raise on this test, but either way sharding shouldn't be silently dropped.
|
Thanks max, I'll get fixes in on this PR and explore how our tests allowed them to slip through |
Fixes the three failures maxrjones identified in review of zarr-developers#4218, whose shared root cause was that from_array is really create_array plus an inverse mapping from stored grid metadata back into the chunks=/shards= parameter space, and that inverse mapping only handled regular grids: - The write_data=True copy path iterated shard regions via the .chunks/.shards accessors, which raise for rectilinear grids. _iter_shard_regions (and _shard_grid_shape) now iterate the stored chunk grid directly, which describes the write regions for every grid kind. - chunks='keep' expanded uniform dimensions of a rectilinear grid to one entry per chunk via write_chunk_sizes, resurrecting the O(nchunks) path and changing resize semantics. The stored chunk_shapes now pass through as-is, preserving the bare-int shorthand in O(ndim). - shards='keep' silently dropped sharding when the shard grid was rectilinear, and .info raised on such arrays through the .shards accessor. The shard grid's chunk_shapes now round-trip through the shards= parameter, and _info guards both accessors by stored grid kind. metadata.chunks now returns the inner chunk shape for any sharded array (inner chunks are always regular, whatever the shard grid), and the repeated sharding-codec isinstance dance is extracted into ArrayV3Metadata.sharding_codec per the existing TODO. ShardsLike admits mixed bare-int/sequence specs, matching what the normalizer accepts. Tests: the ad-hoc from_array regression test is replaced by a spec matrix (FROM_ARRAY_KEEP_CASES) covering every kind of chunk/shard spec create_array accepts, round-tripped through from_array at both write_data settings with grid, codec, and data equality asserted, plus an O(1) pin at shape=(2**62, 30) and error tests for the two accessors that still raise. The stateful hypothesis machine drops its write_data=False workaround, restoring property coverage of the default copy path. Assisted-by: ClaudeCode:claude-fable-5
Follow-up to the from_array 'keep' fixes, addressing the remaining inconsistencies in rectilinear-sharded introspection: - .info / info_complete rendered such arrays as unsharded: no shard line, and the stored-shard count labeled 'Chunks Initialized'. ArrayInfo._shard_shape now admits a '<variable>' sentinel, so the shard shape renders as <variable> and the count is labeled 'Shards Initialized'. - nchunks_initialized raised through the .shards accessor. It now sums the chunk count of each initialized shard individually (shard sizes vary under a rectilinear shard grid; the per-shard sum reduces to the old uniform multiplier for regular shard grids). - The user guide still said .chunks is only available for regular grids; it now documents that sharded arrays always have .chunks, and the rectilinear-shard section shows .chunks/.write_chunk_sizes/.info for such arrays. The round-trip matrix gains an initialized-count assertion, plus tests for the info rendering and the per-shard count. Assisted-by: ClaudeCode:claude-fable-5
|
@maxrjones have a look at the latest changes. one important change: |
|
@ilan-gold do you have time to review this? |
maxrjones
left a comment
There was a problem hiding this comment.
A few more road bumps:
from_array with the default chunks="keep" on a rectilinear source that has been shrunk via resize. Shrinking legally leaves trailing edges beyond the extent (update_shape keeps them on purpose), but the keep path feeds the stored spec into the strict parser, which requires the edges to sum to the span:
def test_from_array_keep_after_shrink_resize() -> None:
src = zarr.create_array(MemoryStore(), shape=(25,), chunks=[[5, 10, 10]], dtype="uint8")
src[:] = np.arange(25, dtype="uint8")
src.resize((18,)) # stored edges legally remain (5, 10, 10)
dst = zarr.from_array(MemoryStore(), data=src, name="0")
assert dst.metadata.chunk_grid == src.metadata.chunk_grid
np.testing.assert_array_equal(dst[:], src[:])Currently raises ValueError: Chunk sizes [5, 10, 10] do not sum to span 18 (_parse_keep_array_attr → normalize_chunks_nd); on main this copy succeeded via the clipped write_chunk_sizes. The shards="keep" branch fails the same way (shards=[[50, 50]], resize((80,))). The normalize_chunks_nd docstring already says stored metadata belongs on the tolerant from_sizes/from_metadata rules — the keep path just needs to follow its own advice (exact preservation as asserted above, or at minimum clipping to the extent like main did). Worth a resize-then-copy entry in FROM_ARRAY_KEEP_CASES.
Out-of-bounds selections on _iter_shard_regions failing loudly, consistent with _iter_shard_keys. The switch from _iter_regions/_iter_grid to ChunkGrid.iter_chunk_regions dropped the bounds check, so the two sibling APIs now disagree on identical inputs:
def test_iter_shard_regions_bounds_check() -> None:
arr = zarr.create_array(MemoryStore(), shape=(30,), chunks=(10,), dtype="uint8")
with pytest.raises(IndexError):
list(arr._iter_shard_keys(origin=(1,), selection_shape=(5,)))
with pytest.raises(IndexError):
list(arr._iter_shard_regions(origin=(1,), selection_shape=(5,)))Currently the second call silently yields 2 regions instead of raising. Since these regions are documented as the safe concurrent-write partitions, a miscomputed selection now silently writes fewer regions rather than failing.
- Mixed per-dimension specs hitting the intended error messages regardless of dimension order. The widened
ChunksLikemakes(5, [10, 20, 70])official, but_is_rectilinear_chunksonly samples the first element, so the v2 and chunks+shards gates are order-dependent:
@pytest.mark.parametrize("chunks", [(5, [10, 20, 70]), ([10, 20, 70], 5)])
def test_mixed_chunks_with_shards_raises(chunks) -> None:
shape = (30, 100) if isinstance(chunks[0], int) else (100, 30)
with pytest.raises(ValueError, match="Rectilinear chunks with sharding is not supported"):
zarr.create_array(MemoryStore(), shape=shape, chunks=chunks, shards=(10, 10), dtype="uint8")
with pytest.raises(ValueError, match="Zarr format 2 does not support rectilinear chunk grids"):
zarr.create_array(MemoryStore(), shape=shape, chunks=chunks, zarr_format=2, dtype="uint8")Currently the bare-int-first parametrization dies later with the internal chunk_shape is only available for regular chunk grids. Use grid[coords] for per-chunk sizes. — advice naming a private API. Checking any(...) across dimensions instead of the first element fixes both gates at once.
Two smaller notes, no test required:
- The changelog line "creation-time counterpart of the indexing fix in #4172" oversells the pairing — #4172 only covers sorted 1-D coordinate selections, and the original #4174 repro still hits the 1.01 TiB
np.bincountallocation atset_coordinate_selectiontime on this branch (creation itself is now instant, as promised). Suggest rewording so #4174 stays open for the indexing half. _nchunks_initializednow walks every shard coordinate even for regular grids, where the oldnshards_initialized * chunks_per_shardmultiply is exact and O(1) after the listing (~4× slower at 10k shards, linear beyond). Keeping the multiply for regular grids and iterating only the initialized keys for rectilinear ones would restore that, and in the spirit of preferring small tests that fail over big ones that hang, an assertion that the two computations agree on a regular sharded grid would pin it.
|
thanks for spotting these! fixes inbound |
ilan-gold
left a comment
There was a problem hiding this comment.
I'm not too familiar with this part of the codebase but it was nice to see it. One thing I don't get even by reading the docs is the difference in {read,write}_chunks_sizes which is pretty relevant here.
https://zarr.readthedocs.io/en/stable/api/zarr/array/#zarr.Array.read_chunk_sizes
https://zarr.readthedocs.io/en/stable/api/zarr/array/#zarr.Array.write_chunk_sizes
These have identical code examples in terms of chunk sizes so I don't fully get it.
| Explicit per-chunk size lists now always produce a rectilinear chunk grid, | ||
| even when the sizes happen to describe a regular grid (all equal, or all equal | ||
| with a smaller trailing chunk). Previously such input was silently collapsed to |
There was a problem hiding this comment.
Just to be clear, this is actually gated behind the feature flag, right?
There was a problem hiding this comment.
yes, but IMO we should consider removing the feature flag and making it on by default soon
… restore iteration bounds checks Addresses maxrjones's second review round on zarr-developers#4218: - from_array with chunks='keep' / shards='keep' failed on sources shrunk by resize: the stored grid legally keeps trailing edges beyond the extent, but the keep path re-validated it through the strict user-spec parser. ChunksLike/ShardsLike now admit a stored ChunkGridMetadata, which normalize_chunks_nd accepts under the tolerant from_sizes rules and init_array stores verbatim — the keep path passes the grid object itself, so the copy preserves the stored grid exactly. This also fixes the last fidelity gap: an externally-written rectilinear grid using bare-int shorthand on every dimension no longer degrades to a regular grid. - _iter_shard_regions silently truncated out-of-bounds selections after the switch to ChunkGrid.iter_chunk_regions, while its sibling _iter_shard_keys raises. all_chunk_coords now applies the same bounds check as _iter_grid, so both APIs fail loudly and consistently. - _is_rectilinear_chunks only sampled the first element, so mixed per-dimension specs with a bare int first bypassed the v2 and chunks+shards gates and died later on an internal error. It now scans all dimensions (and recognizes rectilinear grid metadata). - _nchunks_initialized walked every shard coordinate even for regular shard grids; the O(1) multiply is restored there, and rectilinear shard grids now decode only the initialized keys instead of scanning the full grid. - The zarr-developers#4174 changelog entry no longer claims zarr-developers#4172 as the full indexing counterpart; the coordinate-selection allocation stays tracked in zarr-developers#4174. FROM_ARRAY_KEEP_CASES gains shrunk rectilinear and shrunk rectilinear-sharded rows, plus tests for the bounds check, the order-independent gates, the regular-sharded multiply agreeing with the per-shard sum, and the all-bare-int keep round-trip. Assisted-by: ClaudeCode:claude-fable-5
… into claude/happy-swartz-3813dd
|
@maxrjones the issues you raised should be addressed in the latest commits |
- Contrast read_chunk_sizes and write_chunk_sizes in their docstring examples with a sharded array where the two differ (the previous examples were identical, which is what prompted the question), and route read_chunk_sizes through the _sharding_codec helper. - Validate scalar chunk sizes before handling the -1 sentinel, with an error message that names -1 as valid. - Rename _stored_rectilinear_grid to _stored_rectilinear_grid_or_none. - User guide: state why inner chunks are always regular (the sharding codec requires a uniform inner chunk shape). - Changelog: make explicit that everything in the entry sits behind the experimental rectilinear feature flag. Assisted-by: ClaudeCode:claude-fable-5
Summary
this is a claude-authored fix of ##4174. in
ChunkLayout, instead of materializing all chunk specifications, even regular ones, as numpy array, we specifically model regular chunk grid dimensions as a single shape + extent. This meansChunksTuplecan just disappear, and we can re-use our chunk grid abstractions (FixedDimensionandVaryingDimension), which is nice and simple. This avoids o(num chunks) memory problems and avoids some unnecessary abstractions.this is basically what @maxrjones suggested here: #3899 (comment). My instinct to make the API "uniform" came at the cost of an inefficient representation, so I'm pulling back on the commitment to "uniformity" here.
Here's claude's summary:
Chunk normalization now returns a ChunkGrid whose uniform dimensions are stored as FixedDimension (size + extent) instead of one array entry per chunk, so create_array(shape=(262,), chunks=(1,)) succeeds instantly instead of raising "array is too big", and chunks=(1, 1) on a (231, 2**31) array no longer allocates ~17 GB per dimension.
Explicit per-chunk lists collapse to FixedDimension when they describe a regular grid; genuinely irregular lists become VaryingDimension. Both variants bind chunk sizes to their extent, so the two forms carry the same invariants. The intermediate ChunksTuple type and as_regular_shape helper are removed; create_chunk_grid_metadata consumes the grid directly and serializes uniform dimensions of mixed rectilinear grids as the spec's bare-int step-size shorthand.
Creation-time counterpart of the gh-4174 indexing fix.
Assisted-by: ClaudeCode:claude-fable-5
[Describe what this PR changes and why, in your own words.]
For reviewers
Check if this purported simplification is in fact simpler, and confirm that we don't see any unpleasant performance side effects.
Author attestation
TODO
docs/user-guide/*.mdchanges/